"""
智能解析器 (Intelligent Parser)
基于PyMuPDF的PDF解析器，参考OmniDocBench的解析方式
将PDF转化为结构化的文档元素列表
"""

import fitz  # PyMuPDF
import re
import json
import uuid
from typing import List, Dict, Any, Tuple, Optional
from loguru import logger
from dataclasses import dataclass
import base64
from io import BytesIO
from PIL import Image

@dataclass
class DocumentElement:
    """文档元素数据类"""
    element_id: str
    content: str
    element_type: str
    metadata: Dict[str, Any]


class IntelligentParser:
    """智能解析器类"""
    
    def __init__(self):
        """初始化解析器"""
        self.title_patterns = [
            r'^\s*第[一二三四五六七八九十\d]+章\s*',  # 第X章
            r'^\s*第[一二三四五六七八九十\d]+节\s*',  # 第X节
            r'^\s*\d+\.\s*',  # 1.
            r'^\s*\d+\.\d+\s*',  # 1.1
            r'^\s*\d+\.\d+\.\d+\s*',  # 1.1.1 (三级标题)
            r'^\s*[一二三四五六七八九十]+[、\.]\s*',  # 一、
            r'^\s*\([一二三四五六七八九十\d]+\)\s*',  # (一)
        ]
        
        # 编译正则表达式
        self.compiled_patterns = [re.compile(pattern) for pattern in self.title_patterns]
        
        # 三级标题模式 (如 2.2.1)
        self.level3_pattern = re.compile(r'^\s*(\d+\.\d+\.\d+)\s*(.*)$')
        
    def extract_pdf_content(self, file_path: str) -> List[DocumentElement]:
        """
        从PDF文件中提取内容
        
        Args:
            file_path: PDF文件路径
            
        Returns:
            文档元素列表
        """
        elements = []
        
        try:
            # 打开PDF文档
            doc = fitz.open(file_path)
            logger.info(f"开始解析PDF文件: {file_path}, 共 {len(doc)} 页")
            
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                page_elements = self._extract_page_content(page, page_num, file_path)
                elements.extend(page_elements)
                
            doc.close()
            logger.info(f"PDF解析完成，共提取 {len(elements)} 个元素")
            
        except Exception as e:
            logger.error(f"PDF解析失败: {str(e)}")
            raise
        
        return elements
    
    def _extract_page_content(self, page: fitz.Page, page_num: int, file_path: str) -> List[DocumentElement]:
        """
        从单个页面提取内容
        
        Args:
            page: PyMuPDF页面对象
            page_num: 页码
            file_path: 文件路径
            
        Returns:
            页面元素列表
        """
        elements = []
        
        # 获取页面文本块
        text_blocks = page.get_text("dict")
        
        # 获取页面图像
        image_list = page.get_images()
        
        # 获取页面表格（简单检测）- 暂时注释掉以解决表格过多问题
        # tables = self._detect_tables(page)
        
        # 处理文本块
        text_elements = self._process_text_blocks(text_blocks, page_num, file_path)
        elements.extend(text_elements)
        
        # 处理图像
        image_elements = self._process_images(page, image_list, page_num, file_path)
        elements.extend(image_elements)
        
        # 处理表格 - 暂时注释掉以解决表格过多问题
        # table_elements = self._process_tables(tables, page_num, file_path)
        # elements.extend(table_elements)
        
        return elements
    
    def _process_text_blocks(self, text_blocks: Dict, page_num: int, file_path: str) -> List[DocumentElement]:
        """
        处理文本块
        
        Args:
            text_blocks: PyMuPDF文本块字典
            page_num: 页码
            file_path: 文件路径
            
        Returns:
            文本元素列表
        """
        elements = []
        current_level3_content = ""
        current_level3_title = ""
        current_level3_number = ""
        
        for block in text_blocks.get("blocks", []):
            if "lines" not in block:
                continue
                
            # 提取块中的文本
            block_text = ""
            for line in block["lines"]:
                for span in line["spans"]:
                    block_text += span["text"]
                block_text += "\n"
            
            block_text = block_text.strip()
            if not block_text:
                continue
            
            # 检查是否为三级标题
            level3_match = self.level3_pattern.match(block_text)
            if level3_match:
                # 保存之前的三级标题内容
                if current_level3_content and current_level3_title:
                    element = self._create_text_element(
                        content=current_level3_content.strip(),
                        element_type="NarrativeText",
                        title=current_level3_title,
                        level3_number=current_level3_number,
                        page_num=page_num,
                        file_path=file_path,
                        coordinates=block.get("bbox", [0, 0, 0, 0])
                    )
                    elements.append(element)
                
                # 开始新的三级标题
                current_level3_number = level3_match.group(1)
                current_level3_title = block_text
                current_level3_content = ""
                
                # 创建标题元素
                title_element = self._create_text_element(
                    content=block_text,
                    element_type="Title",
                    title=block_text,
                    level3_number=current_level3_number,
                    page_num=page_num,
                    file_path=file_path,
                    coordinates=block.get("bbox", [0, 0, 0, 0])
                )
                elements.append(title_element)
                
            else:
                # 检查是否为其他类型的标题
                is_title = self._is_title(block_text)
                
                if is_title:
                    # 保存之前的三级标题内容
                    if current_level3_content and current_level3_title:
                        element = self._create_text_element(
                            content=current_level3_content.strip(),
                            element_type="NarrativeText",
                            title=current_level3_title,
                            level3_number=current_level3_number,
                            page_num=page_num,
                            file_path=file_path,
                            coordinates=block.get("bbox", [0, 0, 0, 0])
                        )
                        elements.append(element)
                        current_level3_content = ""
                    
                    # 创建标题元素
                    title_element = self._create_text_element(
                        content=block_text,
                        element_type="Title",
                        title=block_text,
                        level3_number="",
                        page_num=page_num,
                        file_path=file_path,
                        coordinates=block.get("bbox", [0, 0, 0, 0])
                    )
                    elements.append(title_element)
                    
                else:
                    # 普通文本，如果在三级标题下，则累积到当前三级标题内容中
                    if current_level3_title:
                        current_level3_content += block_text + "\n"
                    else:
                        # 独立的文本块
                        element = self._create_text_element(
                            content=block_text,
                            element_type="NarrativeText",
                            title="",
                            level3_number="",
                            page_num=page_num,
                            file_path=file_path,
                            coordinates=block.get("bbox", [0, 0, 0, 0])
                        )
                        elements.append(element)
        
        # 处理最后一个三级标题的内容
        if current_level3_content and current_level3_title:
            element = self._create_text_element(
                content=current_level3_content.strip(),
                element_type="NarrativeText",
                title=current_level3_title,
                level3_number=current_level3_number,
                page_num=page_num,
                file_path=file_path,
                coordinates=[0, 0, 0, 0]
            )
            elements.append(element)
        
        return elements
    
    def _is_title(self, text: str) -> bool:
        """
        判断文本是否为标题
        
        Args:
            text: 文本内容
            
        Returns:
            是否为标题
        """
        text = text.strip()
        if not text:
            return False
        
        # 检查是否匹配标题模式
        for pattern in self.compiled_patterns:
            if pattern.match(text):
                return True
        
        # 检查其他标题特征
        # 1. 文本较短且不以句号结尾
        if len(text) < 100 and not text.endswith('。'):
            # 2. 包含关键词
            title_keywords = ['章', '节', '部分', '概述', '总结', '引言', '结论']
            if any(keyword in text for keyword in title_keywords):
                return True
        
        return False
    
    def _create_text_element(self, content: str, element_type: str, title: str, 
                           level3_number: str, page_num: int, file_path: str, 
                           coordinates: List[float]) -> DocumentElement:
        """
        创建文本元素
        
        Args:
            content: 文本内容
            element_type: 元素类型
            title: 所属标题
            level3_number: 三级标题编号
            page_num: 页码
            file_path: 文件路径
            coordinates: 坐标
            
        Returns:
            文档元素
        """
        element_id = f"text_{uuid.uuid4().hex[:8]}"
        
        metadata = {
            "page_number": page_num + 1,  # 页码从1开始
            "coordinates": coordinates,
            "file_name": file_path.split('/')[-1],
            "element_type": element_type,
            "parent_title": title,
            "level3_number": level3_number,
            "content_length": len(content)
        }
        
        return DocumentElement(
            element_id=element_id,
            content=content,
            element_type=element_type,
            metadata=metadata
        )
    
    def _process_images(self, page: fitz.Page, image_list: List, page_num: int, file_path: str) -> List[DocumentElement]:
        """
        处理图像
        
        Args:
            page: PyMuPDF页面对象
            image_list: 图像列表
            page_num: 页码
            file_path: 文件路径
            
        Returns:
            图像元素列表
        """
        elements = []
        
        for img_index, img in enumerate(image_list):
            try:
                # 获取图像数据
                xref = img[0]
                pix = fitz.Pixmap(page.parent, xref)
                
                if pix.n - pix.alpha < 4:  # 确保不是CMYK
                    # 获取图像尺寸
                    img_width = pix.width
                    img_height = pix.height
                    
                    # 图像过滤条件
                    min_size = 50  # 最小尺寸50像素
                    min_area = 2500  # 最小面积2500像素²
                    min_file_size = 10 * 1024  # 最小文件大小10KB
                    
                    # 转换为PNG格式的base64编码
                    img_data = pix.tobytes("png")
                    img_file_size = len(img_data)
                    img_area = img_width * img_height
                    
                    # 过滤小图像、图标、标点符号等
                    if (img_width >= min_size and 
                        img_height >= min_size and 
                        img_area >= min_area and 
                        img_file_size >= min_file_size):
                        
                        img_base64 = base64.b64encode(img_data).decode()
                        
                        element_id = f"image_{uuid.uuid4().hex[:8]}"
                        
                        # 获取图像在页面中的位置
                        img_rect = img[1:5] if len(img) > 4 else [0, 0, 0, 0]
                        
                        metadata = {
                            "page_number": page_num + 1,
                            "coordinates": img_rect,
                            "file_name": file_path.split('/')[-1],
                            "element_type": "Image",
                            "image_format": "png",
                            "image_size": img_file_size,
                            "image_width": img_width,
                            "image_height": img_height,
                            "image_area": img_area,
                            "image_index": img_index,
                            "quality_score": self._calculate_image_quality_score(img_width, img_height, img_file_size)
                        }
                        
                        element = DocumentElement(
                            element_id=element_id,
                            content=img_base64,
                            element_type="Image",
                            metadata=metadata
                        )
                        elements.append(element)
                        
                        logger.info(f"保存图像: 页面{page_num+1}, 尺寸{img_width}x{img_height}, 大小{img_file_size/1024:.1f}KB")
                    else:
                        logger.debug(f"过滤小图像: 页面{page_num+1}, 尺寸{img_width}x{img_height}, 大小{img_file_size/1024:.1f}KB")
                
                pix = None  # 释放内存
                
            except Exception as e:
                logger.warning(f"处理图像时出错: {str(e)}")
                continue
        
        return elements
    
    def _calculate_image_quality_score(self, width: int, height: int, file_size: int) -> float:
        """
        计算图像质量分数
        
        Args:
            width: 图像宽度
            height: 图像高度
            file_size: 文件大小
            
        Returns:
            质量分数 (0-100)
        """
        # 基于尺寸和文件大小计算质量分数
        area = width * height
        size_score = min(file_size / (100 * 1024), 1.0) * 40  # 文件大小分数 (最高40分)
        area_score = min(area / (500 * 500), 1.0) * 40  # 面积分数 (最高40分)
        ratio_score = min(min(width/height, height/width), 1.0) * 20  # 长宽比分数 (最高20分)
        
        return size_score + area_score + ratio_score
    
    def _detect_tables(self, page: fitz.Page) -> List[Dict]:
        """
        检测表格
        
        Args:
            page: PyMuPDF页面对象
            
        Returns:
            表格信息列表
        """
        tables = []
        
        try:
            # 方法1: 使用PyMuPDF的find_tables功能
            try:
                page_tables = page.find_tables()
                for table in page_tables:
                    table_info = {
                        "bbox": table.bbox,
                        "content": self._extract_table_content(table),
                        "confidence": 0.9,
                        "method": "pymupdf_find_tables",
                        "rows": len(table.extract()),
                        "cols": len(table.extract()[0]) if table.extract() else 0
                    }
                    tables.append(table_info)
                    logger.info(f"检测到表格: {table.bbox}, 行数: {table_info['rows']}, 列数: {table_info['cols']}")
            except Exception as e:
                logger.debug(f"PyMuPDF表格检测失败: {str(e)}")
            
            # 方法2: 基于文本布局的表格检测
            if not tables:
                tables.extend(self._detect_tables_by_layout(page))
            
            # 方法3: 基于特殊字符的表格检测（备用方法）
            if not tables:
                tables.extend(self._detect_tables_by_chars(page))
                
        except Exception as e:
            logger.warning(f"表格检测过程中出错: {str(e)}")
        
        return tables
    
    def _extract_table_content(self, table) -> str:
        """
        提取表格内容
        
        Args:
            table: PyMuPDF表格对象
            
        Returns:
            表格内容字符串
        """
        try:
            table_data = table.extract()
            if not table_data:
                return ""
            
            # 将表格数据转换为文本格式
            content_lines = []
            for row in table_data:
                # 过滤空值并连接单元格
                row_content = [str(cell) if cell is not None else "" for cell in row]
                content_lines.append(" | ".join(row_content))
            
            return "\n".join(content_lines)
        except Exception as e:
            logger.warning(f"提取表格内容时出错: {str(e)}")
            return ""
    
    def _detect_tables_by_layout(self, page: fitz.Page) -> List[Dict]:
        """
        基于文本布局检测表格
        
        Args:
            page: PyMuPDF页面对象
            
        Returns:
            表格信息列表
        """
        tables = []
        
        try:
            # 获取页面的文本块
            text_dict = page.get_text("dict")
            
            # 分析文本块的布局模式
            for block in text_dict.get("blocks", []):
                if "lines" in block:
                    lines = block["lines"]
                    
                    # 检查是否有表格特征
                    if self._is_table_layout(lines):
                        bbox = block["bbox"]
                        content = self._extract_block_text(block)
                        
                        table_info = {
                            "bbox": bbox,
                            "content": content,
                            "confidence": 0.8,
                            "method": "layout_analysis",
                            "lines_count": len(lines)
                        }
                        tables.append(table_info)
                        logger.info(f"基于布局检测到表格: {bbox}")
                        
        except Exception as e:
            logger.warning(f"基于布局的表格检测失败: {str(e)}")
        
        return tables
    
    def _is_table_layout(self, lines: List) -> bool:
        """
        判断文本行是否具有表格布局特征
        
        Args:
            lines: 文本行列表
            
        Returns:
            是否为表格布局
        """
        if len(lines) < 3:  # 至少需要3行才可能是表格
            return False
        
        # 检查行的对齐模式
        x_positions = []
        for line in lines:
            if "spans" in line:
                for span in line["spans"]:
                    x_positions.append(span["bbox"][0])  # 左边界x坐标
        
        if len(x_positions) < 6:  # 需要足够的文本片段
            return False
        
        # 检查是否有规律的列对齐
        x_positions.sort()
        
        # 计算相邻x坐标的差值
        gaps = []
        for i in range(1, len(x_positions)):
            gap = x_positions[i] - x_positions[i-1]
            if gap > 10:  # 忽略小间隙
                gaps.append(gap)
        
        # 如果有多个相似的间隙，可能是表格列
        if len(gaps) >= 2:
            gap_groups = {}
            for gap in gaps:
                # 将相似的间隙分组（容差20像素）
                found_group = False
                for group_gap in gap_groups:
                    if abs(gap - group_gap) <= 20:
                        gap_groups[group_gap] += 1
                        found_group = True
                        break
                if not found_group:
                    gap_groups[gap] = 1
            
            # 如果有重复的间隙模式，可能是表格
            max_count = max(gap_groups.values())
            return max_count >= 2
        
        return False
    
    def _extract_block_text(self, block: Dict) -> str:
        """
        提取文本块的内容
        
        Args:
            block: 文本块字典
            
        Returns:
            文本内容
        """
        text_lines = []
        for line in block.get("lines", []):
            line_text = ""
            for span in line.get("spans", []):
                line_text += span.get("text", "")
            if line_text.strip():
                text_lines.append(line_text.strip())
        
        return "\n".join(text_lines)
    
    def _detect_tables_by_chars(self, page: fitz.Page) -> List[Dict]:
        """
        基于特殊字符检测表格（备用方法）
        
        Args:
            page: PyMuPDF页面对象
            
        Returns:
            表格信息列表
        """
        tables = []
        
        # 获取页面文本
        text = page.get_text()
        
        # 表格特征字符
        table_indicators = [
            '\t',  # 制表符
            '│', '┃',   # 竖线
            '─', '━',   # 横线
            '┌', '┐', '└', '┘',  # 表格角
            '├', '┤', '┬', '┴', '┼',  # 表格连接
            '╔', '╗', '╚', '╝',  # 双线角
            '╠', '╣', '╦', '╩', '╬',  # 双线连接
        ]
        
        # 计算表格字符的密度
        table_char_count = sum(text.count(char) for char in table_indicators)
        total_chars = len(text)
        
        if total_chars > 0:
            table_density = table_char_count / total_chars
            
            # 如果表格字符密度超过阈值，认为可能包含表格
            if table_density > 0.01 or table_char_count > 10:
                table_info = {
                    "bbox": [0, 0, page.rect.width, page.rect.height],
                    "content": text,
                    "confidence": min(0.6, table_density * 10),
                    "method": "character_analysis",
                    "table_char_count": table_char_count,
                    "table_density": table_density
                }
                tables.append(table_info)
                logger.info(f"基于字符检测到表格: 密度{table_density:.3f}, 字符数{table_char_count}")
        
        return tables
    
    def _process_tables(self, tables: List[Dict], page_num: int, file_path: str) -> List[DocumentElement]:
        """
        处理表格
        
        Args:
            tables: 表格列表
            page_num: 页码
            file_path: 文件路径
            
        Returns:
            表格元素列表
        """
        elements = []
        
        for table_index, table in enumerate(tables):
            try:
                element_id = f"table_{uuid.uuid4().hex[:8]}"
                
                # 处理不同格式的bbox
                bbox = table.get("bbox", [0, 0, 0, 0])
                if hasattr(bbox, 'x0'):  # PyMuPDF Rect对象
                    coordinates = [bbox.x0, bbox.y0, bbox.x1, bbox.y1]
                elif isinstance(bbox, (list, tuple)) and len(bbox) >= 4:  # 列表或元组
                    coordinates = list(bbox[:4])
                else:
                    coordinates = [0, 0, 0, 0]
                
                metadata = {
                    "page_number": page_num + 1,
                    "coordinates": coordinates,
                    "file_name": file_path.split('/')[-1],
                    "element_type": "Table",
                    "table_index": table_index,
                    "detection_method": table.get("method", "unknown"),
                    "confidence": table.get("confidence", 0.5),
                    "rows": table.get("rows", 0),
                    "cols": table.get("cols", 0)
                }
                
                element = DocumentElement(
                    element_id=element_id,
                    content=table.get("content", ""),
                    element_type="Table",
                    metadata=metadata
                )
                elements.append(element)
                
                logger.info(f"创建表格元素: 页面{page_num+1}, 方法{table.get('method', 'unknown')}, 置信度{table.get('confidence', 0):.2f}")
                
            except Exception as e:
                logger.warning(f"处理表格 {table_index} 时出错: {str(e)}")
                continue
        
        return elements
    
    def parse_pdf(self, file_path: str) -> List[Dict[str, Any]]:
        """
        解析PDF文件的主要接口
        
        Args:
            file_path: PDF文件路径
            
        Returns:
            文档元素字典列表
        """
        elements = self.extract_pdf_content(file_path)
        
        # 转换为字典格式
        result = []
        for element in elements:
            result.append({
                "element_id": element.element_id,
                "content": element.content,
                "element_type": element.element_type,
                "metadata": element.metadata
            })
        
        return result